package org.springframework.ai.dashscope.metadata;

import org.springframework.ai.dashscope.api.DashscopeAiApi;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.util.Assert;

public class DashscopeAiUsage implements Usage {

	public static DashscopeAiUsage from(DashscopeAiApi.Usage usage) {
		return new DashscopeAiUsage(usage);
	}

	private final DashscopeAiApi.Usage usage;

	protected DashscopeAiUsage(DashscopeAiApi.Usage usage) {
		Assert.notNull(usage, "dashscopeAi Usage must not be null");
		this.usage = usage;
	}

	protected DashscopeAiApi.Usage getUsage() {
		return this.usage;
	}

	@Override
	public Integer getPromptTokens() {
		Integer promptTokens = getUsage().promptTokens();
		return promptTokens != null ? promptTokens : 0;
	}

	@Override
	public Integer getCompletionTokens() {
		return null;
	}

	@Override
	public Integer getTotalTokens() {
		Integer totalTokens = getUsage().totalTokens();
		if (totalTokens != null) {
			return totalTokens;
		}
		else {
			return getPromptTokens() + getCompletionTokens();
		}
	}

	@Override
	public Object getNativeUsage() {
		return null;
	}

	@Override
	public String toString() {
		return getUsage().toString();
	}

}
